import cv2

# 加载人脸识别的分类器
face_cascade = cv2.CascadeClassifier('haarcascade_frontalface_default.xml')

# 加载模型和标签
model = cv2.face.LBPHFaceRecognizer_create()
model.read("trained_model.yml")

# 加载标签
labels = {}
with open("labels.txt", "r") as f:
    lines = f.readlines()
    for line in lines:
        label, name = line.strip().split(":")
        labels[int(label)] = name

# 打开摄像头
cap = cv2.VideoCapture(0)

while True:
    # 读取图像
    ret, img = cap.read()

    # 灰度化
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

    # 检测人脸
    faces = face_cascade.detectMultiScale(gray, scaleFactor=1.5, minNeighbors=5)

    # 识别人脸
    for (x, y, w, h) in faces:
        # 截取人脸区域
        roi_gray = gray[y:y + h, x:x + w]

        # 预处理
        roi_gray = cv2.resize(roi_gray, (100, 100), interpolation=cv2.INTER_LINEAR)

        # 识别
        label, confidence = model.predict(roi_gray)

        # 显示姓名
        if confidence < 50:
            name = labels[label]
            color = (0, 255, 0)
        else:
            name = "Unknown"
            color = (0, 0, 255)

        cv2.putText(img, name, (x, y), cv2.FONT_HERSHEY_PLAIN, 1, color, 2)
        cv2.rectangle(img, (x, y), (x + w, y + h), color, 2)

    # 显示图像
    cv2.imshow("Face Detection", img)

    # 按下q键退出
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

# 释放摄像头
cap.release()

# 关闭所有窗口
cv2.destroyAllWindows()
